import torch
from loguru import logger
from torch import nn


class WaveNetDDoSDetector(nn.Module):
    """基于空洞卷积的WaveNet风格模型"""

    def __init__(self, input_dim: int):
        super().__init__()
        logger.debug("使用空洞卷积+残差连接构建DDoS检测模型")

        # 初始卷积
        self.init_conv = nn.Conv1d(input_dim, 64, kernel_size=1)

        # 空洞卷积块
        self.dilated_convs = nn.ModuleList()
        for i in range(8):
            dilation = 2 ** (i % 5)  # 循环1,2,4,8,16
            self.dilated_convs.append(
                nn.Sequential(
                    nn.Conv1d(64, 64, kernel_size=3,
                              padding=dilation, dilation=dilation),
                    nn.BatchNorm1d(64),
                    nn.GELU(),
                    nn.Dropout(0.1)
                )
            )

        # 全局特征提取
        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self.classifier = nn.Linear(64, 1)

    def forward(self, x):
        # x shape: (batch, features, seq_len)
        x = self.init_conv(x)
        residual = x

        for conv in self.dilated_convs:
            out = conv(x)
            x = out + residual  # 残差连接
            residual = x

        x = self.global_pool(x).squeeze(-1)
        return torch.sigmoid(self.classifier(x))
